package dht

type Router struct {
	local   ID
	k       int
	buckets [160]*Bucket
}

func NewRouter(local ID, k int) *Router {
	r := &Router{}
	r.local = local
	r.k = k
	for i := range r.buckets {
		r.buckets[i] = &Bucket{}
	}
	return r
}

func (r *Router) bucketIndex(id ID) int {
	var bi int
	if r.local == id {
		bi = len(r.buckets) - 1
	} else {
		bi = len(r.buckets) - r.local.Xor(id).BitLen()
	}
	return bi
}

func (r *Router) Len() int {
	length := 0
	for _, bucket := range r.buckets {
		length = length + bucket.Len()
	}
	return length
}

func (r *Router) Update(id ID) {
	if r.local == id {
		return
	}
	bi := r.bucketIndex(id)
	bucket := r.buckets[bi]
	bucket.Put(id, r.k)
}

func (r *Router) Remove(id ID) {
	bi := r.bucketIndex(id)
	bucket := r.buckets[bi]
	bucket.Del(id)
}

func (r *Router) Near(target ID, k int) IDList {
	ret := make(IDList, 0, k)
	bi := r.bucketIndex(target)
	for i := bi; i < len(r.buckets); i++ {
		if len(ret) >= k {
			break
		}
		bucket := r.buckets[i]
		ret = append(ret, bucket.List()...)
	}
	for i := bi - 1; i >= 0; i-- {
		if len(ret) >= k {
			break
		}
		bucket := r.buckets[i]
		ret = append(ret, bucket.List()...)
	}
	ret = ret.Sort(target)
	if len(ret) > k {
		ret = ret[:k]
	}
	return ret
}

func (r *Router) Each(fn func(id ID) bool) bool {
	for _, bucket := range r.buckets {
		flag := bucket.Each(fn)
		if flag == false {
			return false
		}
	}
	return true
}
